package com.loklok.code.confuse.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import okhttp3.*;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class AsyncOpenAIUtil {

    public static final ExecutorService OPENAI_EXECUTOR = Executors.newFixedThreadPool(
            Math.min(2, Runtime.getRuntime().availableProcessors()));

    static {
        Runtime.getRuntime().addShutdownHook(new Thread(OPENAI_EXECUTOR::shutdown));
    }


    private static final String regex = "```([\\s\\S]*?)```";
    private static final Pattern pattern = Pattern.compile(regex);

    private static final OkHttpClient CLIENT = new OkHttpClient.Builder()
            .connectTimeout(2, TimeUnit.MINUTES)
            .readTimeout(2, TimeUnit.MINUTES)
            .writeTimeout(2, TimeUnit.MINUTES)
            .build();

    // 使用回调接口替代返回值
    public interface ResultCallback<T> {
        void onSuccess(T result);
        void onFailure(Exception e);
    }

    public static void getParamBodyAsync(String apiKey, String prompt, ResultCallback<String> callback) {
        doDeepSeekRequestAsync(apiKey, prompt, new ResultCallback<String>() {
            @Override
            public void onSuccess(String content) {
                // 正则表达式匹配 ``` 包裹的代码块
                Matcher matcher = pattern.matcher(content);
                // 查找并输出匹配的代码块
                if (matcher.find()) {
                    String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
                    content = codeBlock.replace("java\n", "");
                }
                callback.onSuccess(content);
            }

            @Override
            public void onFailure(Exception e) {
                callback.onFailure(e);
            }
        });
    }

    public static void getSynonym(String apiKey, String prompt, ResultCallback<Map<String,Set<String>>> callback){


        doDeepSeekRequestAsync(apiKey, prompt, new ResultCallback<String>() {
            @Override
            public void onSuccess(String content) {
                Matcher matcher = pattern.matcher(content);
                // 查找并输出匹配的代码块
                if (matcher.find()) {
                    String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
                    content = codeBlock.replace("json\n", "");
                }
                Map<String, Set<String>> map = JSON.parseObject(content,  new TypeReference<Map<String, Set<String>>>() {}.getType());
                callback.onSuccess(map);
            }

            @Override
            public void onFailure(Exception e) {
                callback.onFailure(e);
            }
        });

    }

    public static void getVmBodyAsync(String apiKey, String prompt, ResultCallback<String> callback) {
        doDeepSeekRequestAsync(apiKey, prompt, new ResultCallback<String>() {
            @Override
            public void onSuccess(String content) {
                Matcher matcher = pattern.matcher(content);
                // 查找并输出匹配的代码块
                if (matcher.find()) {
                    String codeBlock = matcher.group(1); // 获取第一个捕获组（即代码内容）
                    content = codeBlock.replace("vm\n", "");
                }
                callback.onSuccess(content);
            }

            @Override
            public void onFailure(Exception e) {
                callback.onFailure(e);
            }
        });
    }

    private static void doDeepSeekRequestAsync(String apiKey, String content, ResultCallback<String> callback) {
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("model", "deepseek-chat");
        jsonObject.put("stream", false);

        JSONArray jsonArray = new JSONArray();
        JSONObject messageObject = new JSONObject();
        messageObject.put("role", "user");
        messageObject.put("content", content.toString());
        jsonArray.add(messageObject);
        jsonObject.put("messages", jsonArray);

        RequestBody body = RequestBody.create(
                jsonObject.toJSONString().getBytes(StandardCharsets.UTF_8),
                MediaType.get("application/json; charset=utf-8")
        );

        Request request = new Request.Builder()
                .url("https://api.deepseek.com/chat/completions")
                .post(body)
                .addHeader("Authorization", "Bearer " + apiKey)
                .addHeader("Content-Type", "application/json")
                .build();

        // 使用 OPENAI_EXECUTOR 提交任务
        OPENAI_EXECUTOR.execute(() -> {

            CLIENT.newCall(request).enqueue(new okhttp3.Callback() {
                @Override
                public void onFailure(Call call, IOException e) {
                    callback.onFailure(e);
                }

                @Override
                public void onResponse(Call call, Response response) throws IOException {
                    if (!response.isSuccessful()) {
                        callback.onFailure(new IOException("Unexpected code " + response));
                        return;
                    }

                    String res = response.body().string();
                    JSONObject resJSONObject = JSON.parseObject(res);
                    JSONObject choicesOne = resJSONObject.getJSONArray("choices").getJSONObject(0);

                    if (!"stop".equals(choicesOne.getString("finish_reason"))) {
                        callback.onFailure(new RuntimeException("deep seek 回答未结束"));
                        return;
                    }

                    String resContent = choicesOne.getJSONObject("message").getString("content");
                    callback.onSuccess(resContent);
                }
            });

        });
    }
}



